# 验证过拟合
# 分类任务可以这么做，生成式任务不能这么做
# 所有的文本分类都可以这么做
import os
import torch
from torch.utils.data import Dataset, DataLoader
from datasets import load_dataset, load_from_disk
from transformers import BertModel, BertTokenizer
from torch.optim import AdamW
import tensorflow as tf

os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"

model_path = r"D:\code\base_learn\learn_ai\hugging_face\model\bert-base-chinese\models--bert-base-chinese\snapshots\c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f"

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 定义设备信息
pretrained_model = BertModel.from_pretrained(model_path).to(DEVICE)
# 定义训练的轮次（将整个数据集训练完一次为一轮）
EPOCH = 10


# 2_my_data.py
class MyDataset(Dataset):
    def __init__(self, split):
        """
        初始化数据集
        :param split:
        """
        # 从磁盘加载数据
        self.dataset = load_from_disk(r"D:\code\base_learn\learn_ai\hugging_face\saved_dataset")
        if split == "train":
            self.dataset = self.dataset["train"]
        elif split == "test":
            self.dataset = self.dataset["test"]
        elif split == "validation":
            self.dataset = self.dataset["validation"]
        else:
            raise ValueError("Invalid split")

    def __len__(self):
        """
        返回数据集长度
        :return:
        """
        return len(self.dataset)

    def __getitem__(self, index):
        """
        对每条数据单独做处理
        :param index:
        :return:
        """
        my_text = self.dataset[index]["text"]
        my_label = self.dataset[index]["label"]
        return my_text, my_label


# 3_net.py
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 设计全连接网络，实现二分类任务
        self.fc = torch.nn.Linear(768, 2)

    # 使用模型处理数据（执行前向计算）
    def forward(self, input_ids, attention_mask, token_type_ids):
        # 冻结 Bert 模型的参数，让其不参与训练
        with torch.no_grad():
            outputs = pretrained_model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                token_type_ids=token_type_ids
            )
        # 增量模型参与训练
        outputs = self.fc(outputs.last_hidden_state[:, 0])
        return outputs


# 加载字典和分词器
token = BertTokenizer.from_pretrained(model_path)


# 将传入的字符串进行编码
def collect_fn(data):
    """
    对加载的数据进行编码
    :param data:
    :return:
    """
    sentences = [i[0] for i in data]
    label = [i[1] for i in data]
    # 对文本进行编码
    encodings = token.batch_encode_plus(
        batch_text_or_text_pairs=sentences,
        add_special_tokens=True,  # 是否添加 [CLS] 和 [SEP] 特殊字符
        truncation=True,  # 当句子长度大于 max_length（上限是 model_max_length） 时，截断句子
        max_length=512,  # 句子最大长度
        padding="max_length",  # 填充方式，一律补 0 到 max_length
        return_tensors="pt",  # 可取值 "pt" 或 "tf"， 默认为 list
        # 返回序列长度
        return_length=True,
    )
    return encodings["input_ids"], encodings["attention_mask"], encodings["token_type_ids"], torch.LongTensor(label)


# 创建数据集
train_dataset = MyDataset("train")
train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=50,  # 训练批次
    shuffle=True,  # 打乱数据集
    drop_last=True,  # 舍弃最后一个批次的数据，防止形状出错
    collate_fn=collect_fn  # 对加载的数据进行编码
)
# 验证数据集
val_dataset = MyDataset("validation")
val_loader = DataLoader(
    dataset=train_dataset,
    batch_size=50,  # 训练批次
    shuffle=True,  # 打乱数据集
    drop_last=True,  # 舍弃最后一个批次的数据，防止形状出错
    collate_fn=collect_fn  # 对加载的数据进行编码
)

if __name__ == '__main__':
    # 开始训练
    print(DEVICE)
    model = Model().to(DEVICE)
    # 定义优化器
    optimizer = AdamW(model.parameters())
    # 定义损失函数
    loss_fn = torch.nn.CrossEntropyLoss()

    # 初始化验证最佳准确率
    best_val_acc = 0.0

    # 开始训练
    for epoch in range(EPOCH):
        for i, (input_ids, attention_mask, token_type_ids, label) in enumerate(train_loader):
            # 将数据放到 DEVICE 上
            input_ids = input_ids.to(DEVICE)
            attention_mask = attention_mask.to(DEVICE)
            token_type_ids = token_type_ids.to(DEVICE)
            label = label.to(DEVICE)
            # 前向计算
            outputs = model(input_ids, attention_mask, token_type_ids)
            # ################# 使用 pytorch
            # 根据输出计算损失
            loss = loss_fn(outputs, label)
            # 根据误差优化参数
            optimizer.zero_grad()  # 清空梯度
            loss.backward()  # 反向传播，计算梯度，自动求导
            optimizer.step()   # 优化参数，梯度下降

            # 每隔五个批次输出训练信息
            if i % 5 == 0:
                out = torch.argmax(outputs, dim=1)
                # 计算训练精度
                acc = (out == label).sum().item() / label.shape[0]
                print(f"epoch: {epoch}, batch: {i}, acc: {acc}, loss:{loss.item()}")

        # 验证模型，判断模型是否过拟合
        # 设置为评估模式
        model.eval()
        # 不需要模型参与训练
        with torch.no_grad():
            val_acc = 0.0   # 验证精度
            val_loss = 0.0  # 验证损失
            for i, (input_ids, attention_mask, token_type_ids, label) in enumerate(val_loader):
                # 将数据放到 DEVICE 上
                input_ids = input_ids.to(DEVICE)
                attention_mask = attention_mask.to(DEVICE)
                token_type_ids = token_type_ids.to(DEVICE)
                label = label.to(DEVICE)
                # 前向计算
                outputs = model(input_ids, attention_mask, token_type_ids)
                # ################# 使用 pytorch
                # 根据输出计算损失
                val_loss = loss_fn(outputs, label)
                # 根据数据，计算验证精度
                out = torch.argmax(outputs, dim=1)
                val_acc += (out == label).sum().item()
            val_loss /= len(val_loader)
            val_acc /= len(val_loader)
            print(f"验证集: val_acc: {val_acc}, val_loss:{val_loss.item()}")
            #  根据验证准确率保存最优参数
            if val_acc > best_val_acc:
                best_val_acc = val_acc
                torch.save(model.state_dict(), f"params1/best_bert.pth")
                print(f"最佳模型模型成功，保存路径为：params1/best_bert.pth")
        #  保存最后一轮参数
        torch.save(model.state_dict(), f"params1/last_bert.pth")
        print(f"保存模型成功，保存路径为：params1/last_bert.pth")